"""
 Copyright (c) 2022 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
      http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

from typing import List, Optional, Union
import numpy as np

from nncf.common.quantization.structs import QuantizationMode
from nncf.common.quantization.structs import QuantizerConfig
from nncf.common.tensor_statistics.statistics import MinMaxTensorStatistic


class ONNXQuantizerLayerParameters:
    """
    Class handles Quantizer/Dequantizer layer attributes.
    """

    def __init__(self, scale: List[float], zero_point: List[int], mode: QuantizationMode, axis: Optional[int]):
        self.scale = scale
        self.zero_point = zero_point
        self.mode = mode
        self.axis = axis


def calculate_scale_level(max_val: Union[float, np.ndarray], min_val: Union[float, np.ndarray],
                          num_bits: int,
                          mode: QuantizationMode) -> Union[float, np.ndarray]:
    """
    Calculates Quantizer/Dequantizer layer scale level.
    """
    if mode == QuantizationMode.SYMMETRIC:
        input_abs_max = np.maximum(np.abs(max_val), np.abs(min_val))
        return input_abs_max / ((2 ** num_bits - 1) / 2)
    return (max_val - min_val) / 2 ** num_bits


def calculate_weight_quantizer_parameters(weight_tensor: np.ndarray, quantizer_config: QuantizerConfig,
                                          axis: Optional[int]) -> ONNXQuantizerLayerParameters:
    """
    Calculates Quantizer/Dequantizer layer attributes for weight quantizer such as scale, zero_points and
    quantization mode: symmetric, asymmetric.

    :param weight_tensor: Weight tensor to calculate quantizer attributes.
    :param quantizer_config: Config of Quantizer.
    :param axis: In per-channel case - the axis for the quantization. In per-tensor - ignored.
    :return: Parameters of Quantizer.
    """
    per_channel = quantizer_config.per_channel
    num_bits = quantizer_config.num_bits
    mode = quantizer_config.mode

    if per_channel:
        assert axis is not None
        axes = list(range(len(weight_tensor.shape)))
        axes.pop(axis)
        axes = tuple(axes)
    else:
        axes = None
    input_high = np.amax(weight_tensor, axis=axes)
    input_low = np.amin(weight_tensor, axis=axes)
    scales = calculate_scale_level(input_high, input_low, num_bits, mode)
    zero_points = np.zeros_like(scales, dtype=np.int64)
    return ONNXQuantizerLayerParameters(scales.tolist(), zero_points.tolist(), mode, axis)


def calculate_activation_quantizer_parameters(statistics: MinMaxTensorStatistic,
                                              quantizer_config: QuantizerConfig,
                                              axis: Optional[int] = None) -> ONNXQuantizerLayerParameters:
    """
    Calculates Quantizer/Dequantizer layer attributes for activation quantizer such as scale, zero_points and
    quantization mode: symmetric, asymmetric.

    :param statistics: Collected statistics for the quantized insertion.
    :param quantizer_config: Config of the quantization configuration.
    :param axis: Axis of the quantization. None in a per-tensor quantization case.
    :return: Parameters of the quantizer/dequantizer layers.
    """
    per_channel = quantizer_config.per_channel
    num_bits = quantizer_config.num_bits
    mode = quantizer_config.mode
    input_low = statistics.min_values
    input_high = statistics.max_values
    if per_channel:
        assert axis is not None
        raise RuntimeError('Currently per-channel is not supported for activation tensors.')
    scales = calculate_scale_level(input_high, input_low, num_bits, mode)
    zero_points = np.zeros_like(scales, dtype=np.int64)
    return ONNXQuantizerLayerParameters(scales.tolist(), zero_points.tolist(), mode, axis)
